Dynamic shape on the C++ trace path (elementwise add)#269
Open
YWHyuk wants to merge 46 commits into
Open
Conversation
c2a243b to
6b39c94
Compare
b7c1ec4 to
9033945
Compare
8e9db1a to
82e9255
Compare
c166abd to
ed5c747
Compare
… feed Skeleton + EmitC + cost/dep analysis on the frontend; the trace runtime, loader, bridge, and Core feed on the simulator; shared MLIR pass helpers and the pipeline tests.
Per-record tag key in the bridge plus per-iteration tag alloc in dma-fine-grained so multi-tile-K and conv loads do not collide; strip the reduction accum marker from the memory_barrier slot.
togsim_dispatch with TILE_BEGIN/TILE_END; outline each work-item into togsim_kernel_tile.
DMA-capacity throttle and frozen-state guard, per-core VMEM in the configs, and the SA weight-buffer throttle.
trace_timeline.py with per-work-item grouping and resource-centric DMA lanes; the trace logs the first DRAM response and the assigned systolic array, and scopes the compute barrier to its dispatch.
Default to the trace path; fix uninitialized Instruction fields, the matmul accumulator wedge, fused-subtile dedup, nested/fused epilogue dataflow, and dma_wait fusion; bound concurrent dispatches to the spad, round-robin work-items within a partition, benchmark autotune and run the multi-tenant scheduler through the trace path, and emit trace.so for pooling/reduction.
Carry simulator headers through the wrapper for cache-safe replay; drop verbose [P3-trace] logs; fix the key.mlir compile race in load().
… runtime model Replace the trace bridge's accumulated special cases with one dataflow rule and clean up the runtime that consumes it. Dependency rule: per SRAM buffer keep a writers SET; a reader depends on all current writers (occupancy=ISSUE when both are systolic-array ops, else latency=DONE); a writer REPLACEs the set. The only exception is is_mm_accum (a matmul that reads and writes the same buffer = a commutative accumulator): skip its read edge and UNION its write, waiting only the non-matmul init seed and not ordering co-matmuls. This drops the matmul-accumulator chain that deadlocked the SA weight-slot pipeline while keeping the init->matmul edge, and lets a vector epilogue or the store wait every K matmul (fixes the pure-vector store that an empty COMPUTE_BAR let slip). Remove COMPUTE_BAR entirely: a matmul is its own DONE-handle (finish == SA drain), so the store JOINs the matmul writers directly. The whole emit/loader chain is gone -- build_skeleton, lower_to_emitc, togsim.compute_barrier, the runtime symbol, the Opcode/case/_fence_finish, and TraceRec::COMPUTE_BAR -- so a stale producer fails loudly instead of emitting records the bridge would drop. Only MEMORY_BAR remains (an async load's DONE is its data arrival, not issue). Model compute-output spad footprint in the SRAM version/capacity machinery so buffer reuse (WAR) is capacity-modeled, not a hard edge. The output size comes from the DMA records that touch the same buffer (a buf_bytes pre-pass); an in-place buffer (accumulator, relu) is version-transparent so footprint is not double-counted. The occupy gate and version release sit in the MOVIN/MOVOUT/COMP issue points (release before the COMP skip path so a skipped matmul still frees). Runtime: collapse child_inst / _pipeline_children into one event-indexed _deps[ISSUE|DONE] with add_dep(c, on) and fire(e); collapse the weight-slot release queue and the async-load wakeup into one _due_events timed-effect table drained by process_due_events. Both are behavior-preserving (byte-identical). Require the weight-slot model: sa_weight_buffer_depth must be > 0 (errors at init), and the round-robin disable mode is removed. Degenerate traces (a consumer-less preload, an unpinned matmul) hit explicit error+exit guards rather than asserts that vanish under NDEBUG. Mark the legacy ONNX TOG path deprecated: it is superseded by the trace path, so TileGraphParser logs a deprecation warning and the TORCHSIM_LEGACY_TOG=1 opt-in warns at command build.
… spad/2
The validation-binary spad-overflow check sat inside `if functional_mode:`, so in
timing-only / autotune (non-functional) runs an over-spad tiling was never
rejected and reached TOGSim, which wedged ("spad too small") and crashed the
compile via assert 0. Move the compile + check out of the functional gate (the
Spike execution itself stays gated, run_spike below) and budget per dispatch at
spad/2 -- two work-items run concurrently (double buffer), so each must fit half
the spad or they deadlock competing for it. This matches the GEMM tiling gate
(max_spad_size = spad/2), which pointwise ops lacked. Fixes the resnet /
test_scheduler wedge where a fused BatchNorm+ReLU tile exceeded spad/2.
Every generated trace.cpp now opens with a "GENERATED ... DO NOT EDIT" banner carrying the TOGSim ABI version and the togsim_* call formats, so a dumped trace is self-documenting. Both are read from togsim_runtime.h at codegen time -- the version from the TOGSIM_ABI_VERSION define, the call-format text from a marked block next to the declarations -- so they never drift when the ABI changes.
…rint The bridge sums each work-item's distinct-buffer footprint (each buffer once, so a reduction's reloads of the same section do not inflate it) onto its Tile. can_issue then admits two concurrent dispatches only when each fits half the spad; a dispatch larger than spad/2 runs alone with the whole spad, so two work-items never compete for the shared spad and deadlock -- a runtime safety net beneath the codegen spad/2 gate. The footprint and resulting max_dispatch are logged on the TILE_SCHEDULED trace line for debugging.
…election select_tile fed the epilogue buffer count to gemm_combination_mapping as max(n_extra_read - 2, 0), which optimistically assumed the X and W DMA buffers were already freed and reusable by the epilogue. The codegen lays every buffer out as a disjoint .spad global and never reuses freed space, so the estimate undercounted the real footprint: for a fused matmul+relu kernel the budget came to 64512 B/lane (treated as a bare GEMM) while the emitted kernel used 89088 B/lane including the relu output buffer. Under the full-spad guard this was harmless (89088 < 131072), but the spad/2 guard rejected it and crashed test_transformer_fusion with SpadOverflowError, since the heuristic path has no tile-shrink fallback. Pass n_extra_node + n_extra_read instead: one output-tile-sized buffer per epilogue node plus one per extra read operand, matching what the codegen emits. For the matmul+relu kernel the budget now equals the actual footprint exactly, and tile selection picks TILE_M=128 (62976 B/lane) which fits spad/2. Liveness-based spad reuse and broadcast over-allocation are tracked separately as optimizations in issue #275.
The spad-overflow check summed the kernel stack frame into the per-lane scratchpad usage (spad_usage = stack_size + spad_size) and rejected the tiling when that exceeded the spad/2 budget. But only the .spad section actually lives in the scratchpad -- it is pinned there by the --section-start=.spad link option. The kernel stack is in main memory (sp is set up by pk in the -m region, not at the scratchpad vaddr), so it does not consume the scratchpad and must not be charged against it. The scalar frame is also shared across lanes, not per-lane, so adding it double-counted. On small configs (8x8) this falsely rejected feasible tilings: the wrapper3 conv2d/resnet/mistral kernels fit the 32 KB spad with room to spare but were tripped over the spad/2 gate purely by the ~200-800 B stack term, crashing compile with SpadOverflowError. Drop the stack term; the .spad-only check still correctly rejects real buffer overflows (e.g. sparsity, which is fixed separately by f05ac8a).
Spike v1.0.3 zero-inits the MVIN/MVOUT DMA address buffer so ROUNDUP padding entries are skipped instead of dereferencing uninitialized garbage, fixing the host store segfault on 8-lane (wrapper3) configs where a tile's split axis is not a multiple of vlane_stride*n_vu.
Add pytorchsim_functional_verify_per_kernel, a sub-option of pytorchsim_functional_mode that localizes the first compiled kernel whose Spike output diverges from a CPU reference. When enabled, the generated wrapper compares every realized buffer (the output of each fused kernel) against a CPU "golden" computed once per call by running the original aten graph (V.graph.module) on CPU with the same inputs, via an fx.Interpreter that records each node's output by name. A buffer is mapped to its originating fx node through V.graph.get_buffer().origin_node, so the check reports the kernel, the originating op, the offending indices and the max abs diff, then raises at the first divergence (stop-at-first). Comparison granularity is the fused-cluster output, the finest observable in a fused pipeline. Auto-disabled when functional mode is off (no Spike values to verify); the config accessor AND-gates the key with functional mode. Codegen bakes the verify_init/verify_check calls into the wrapper only when enabled at compile time, so clear the codegen cache when toggling. Tolerances via TORCHSIM_FUNCTIONAL_VERIFY_RTOL / _ATOL (default 1e-4). extension_functional_verify.py holds the graph registry and the runtime golden/compare logic; mlir_codegen_backend injects the calls at the wrapper level; extension_config reads the new YAML key.
The BMM tile selector fed only n_extra_node to gemm_combination_mapping and dropped the prologue's extra-read operands entirely, so a softmax-fused attention matmul (value^T @ softmax(scores)) was tiled as a bare BMM. The codegen lays the softmax max/sum operands out as their own disjoint weight-tile-sized .spad globals (buf3/buf4) and never reuses freed space, so the estimate undercounted the real footprint: on the 32x32 wrapper2 config (16 KB/lane spad/2 budget) tile selection picked TILE_N=512, whose emitted kernel used 16896 B/lane and overflowed the budget by 512 B, crashing test_transformer with SpadOverflowError. Under the 128x128 full-budget config the slack hid it. This mirrors the epilogue n_extra_read gap fixed for the GEMM template (commit f05ac8a), now on the BMM prologue path. Add an n_prologue_extra_read knob to gemm_combination_mapping that charges each extra prologue-read operand as one weight-tile-sized (TILE_K x TILE_N) .spad buffer, matching what the codegen emits, and have the BMM template count those operands the same way codegen does (the one numel-matching main input reuses the matmul-operand buffer; every other read gets its own global). Tile selection now rejects TILE_N=512 (16896 B/lane) and picks a fitting tile (8704 B/lane), so wrapper2 test_transformer passes the full Spike + allclose check. The new parameter defaults to 0, leaving the GEMM and conv callers unchanged.
Add docs/per-kernel-functional-verify.md (usage, options, mechanism, limitations, code map) and point to it from CLAUDE.md: a one-line entry in the debugging section and in the YAML knobs list.
Keep only a one-line note on what TOGSIM_ABI_VERSION guards; the per-bump v1..v12 history was noise nobody reads.
[Frontend] Budget fused-prologue spad buffers in BMM tile selection
…s off format_dma_inst_issued_trace_line and format_instruction_detail_line are only ever used as the message argument to trace_instruction_line (a spdlog::trace sink). Because the argument is evaluated eagerly at the call site, the kernel builds a formatted std::string for every issued/finished instruction and then spdlog drops it whenever the level is above trace (the default). Bail out of both formatters when trace logging is disabled so the fmt::format work is skipped entirely. This removes wasted work but is not the simulation-speed bottleneck (that is the per-cycle ready-queue rescan in Core::cycle, addressed separately); the conv kernel wall time is unchanged within noise.
Core::cycle() walked every tile's ready-instruction list on every simulated cycle to find one instruction to issue. Profiling an 8x8 conv kernel showed this scan is ~96% of simulation time and, within it, the walk is >99%: the ready list holds ~2000 instructions (mostly blocked on the SRAM-capacity / weight-slot throttle), only one issues per cycle, and the blocked ones are never removed, so the same ~2000 are re-examined every cycle (~2.6k list iterations/cycle, ~1e9 over 400k cycles). At small arrays this dominates because tiny tiles inflate both the ready-list length and the number of DMA-wait stall cycles. Gate the scan behind a per-Core _issue_dirty flag. The scan's outcome can only change when the ready set grows or a resource frees, so set the flag on exactly those events: - ready set grows: Instruction::dec_ready_counter, on enqueueing a now-ready instruction, sets the OWNING core's _issue_dirty via _owner_dirty_ref (wired in Core::issue when the tile is dispatched) -- per-core so a remote core's enqueue does not force every core to rescan; - SRAM freed (release_sram) and weight slot freed (apply_due) set _issue_dirty; - a new dispatch (issue) and a successful issue keep it set. On a cycle with none of these the scan would re-walk the same blocked instructions and issue nothing, so it is skipped. Instructions are never moved between queues, so there is no re-admission churn. Issue-identical: it never changes which instruction issues or when, only whether the (side-effect-free under those conditions) scan runs. The 8x8 conv kernel produces the same 403026 cycles as before, with its TOGSim wall dropping from 214.5s to 41.9s (~5x, 71% of scans skipped). A forced-scan invariant check (no skipped cycle ever issues or inline-finishes a zero-cycle COMP) found 0 violations across 39 kernels spanning GEMM, BMM, softmax and conv. DMA-bound kernels (the ones that hit the 6h CI cap) stall more and gain more.
In Core::cycle()'s issue scan, a COMP with compute_cycle == 0 finishes inline and calls instructions.erase(it), but does not set `issued` or break -- so control falls through to the `it++` at the end of the loop body, incrementing the std::list iterator that erase() just invalidated. That reads a freed node and is undefined behavior; it usually limps along because the freed node's next pointer is briefly intact, but under a different allocator / sanitizer / build it can crash or skip-or-duplicate the remaining ready instructions in that tile, corrupting issue order. Use the iterator erase() returns and continue, the standard erase-in-loop idiom. No behavior change on the path that happened to work; it just removes the UB.
…X TOG The main compile/sim path no longer generates or selects the legacy ONNX Tile-Operation-Graph. extension_codecache emits only trace.so + trace_cycles.tsv (the build-skip now keys on trace.so), and TOGSimulator.run_standalone always drives TOGSim with --trace_so. The TORCHSIM_LEGACY_TOG opt-in is removed from the frontend. The ONNX --models_list branch is kept solely for the STONNE sparse path (extension_op.py); TOGSim's C++ ONNX parser is untouched (separate PR). origins (which FX nodes a kernel came from) is preserved: logged per kernel run and recorded as a trailing "# origins:" line in trace_cycles.tsv -- the legacy ONNX TOG carried this as node metadata, and the C++ cycle-table loader stops at the comment so the current parser is unaffected. Also drop the dead tog_file param from mlir_gem5_compile_command, migrate scripts/chiplet.sh to --trace_so/--cycle_table (the trace path stubs per-tensor addresses and --attributes_list is no longer a Simulator option), and refresh the CLAUDE.md TOG-generation notes.
Rework per-kernel codegen so the loop body is an ordered list of load->compute->store steps (class Step + push_step + codegen_loops iteration) instead of the fixed dma_loads/compute/dma_stores buffers. A step may be compute-only or transfer-only; the compute_idx loop is emitted only for steps that actually have compute content. Drop the ad-hoc self.masks buffer: the reduction-tail mask (get_mask) now emits into the compute buffer, since a mask is just compute. Use the step model to express indirect (gather/scatter) access cleanly. When the indirect index is produced in the same pass, convert_indirect_indexing pushes a new step for the offset build, so the index load, the offset build and the gather become separate steps bridged through spad. This replaces the compute_dependecy -> dma_stores hack in both convert_indirect_indexing and load(). push_step clones the CSE so the name counter is shared but the dedup cache is per-step (each step is its own compute_idx region). Validated: x[idx+2]+1 emits two steps and matches CPU; add/matmul/reduce/softmax/layernorm and gather/scatter unchanged.
…ptor
Replace the affine.apply{indirect_access} symbol smuggle with an explicit
offset descriptor. convert_indirect_indexing returns the offset spad instead
of folding sympy.Symbol(out) into the index; emit_transfer carries it as a
togsim.transfer operand; decompose_transfer lifts that operand to a
memref.dma_start {indirect_offset = @spad_symbol} attribute (memref.dma_start
is a registered op and rejects an extra operand, but accepts the attribute);
lower_dma_to_gemmini reads the attribute and resolves the global for CONFIG4
(drops _find_indirect); build_skeleton adds the offset spad to the gather DMA
read_bufs so the offset-build -> gather dependency edge forms in the trace.
The index stays clean (base only). Validated on both paths: Spike functional
(computed-index gather + scatter allclose, pointwise/reduce regression 0) and
the C++ trace timing path end-to-end (allclose; gather togsim_dma read_bufs now
includes the offset spad).
Indirect addressing (scattered-DMA timing) in the new trace path is a separate
gap tracked in issue #284.
…ride, compute-loop multi-dim offset Three refinements on top of the explicit offset descriptor: 1. Detect indirect access by an explicit symbol set instead of an "tmp" substring match. indirect_indexing now records str(index_var) in self.indirect_symbols; a _has_indirect(expr) helper tests the index free_symbols against it; the former "tmp"-string sites (store, get_dma_info, convert entry/indirect_dims/stride) use it. Removes the now-dead _find_indirect from lower_dma_to_gemmini. 2. Single indirect dim: pass the raw index spad and let the MVIN apply the gather stride per position (CONFIG4 offset_stride) instead of a vector_load/muli/vector_store round-trip that pre-multiplied the stride. emit_transfer carries offset_stride; decompose copies the attribute; lower programs CONFIG4 with it (was hardcoded to 1). 3. Multi indirect dim (e.g. x[ix, iy]): the offset is a sum of strided indices, which a single CONFIG4 channel cannot do, so the sum stays in the kernel -- but build it in the compute loop (chunked by compute_vec_size, not a single tile-wide vector) and store it to a DEDICATED offset spad so an index that is live elsewhere (x[ix, iy] + ix) is not clobbered. push_step separates the offset-build loop from the gather that reads it. Adds multi-dim gather and index-reuse cases to test_indirect_access. Validated: indirect/scatter/embedding + the two new multi-dim cases pass; add/matmul/softmax regression 0.
…emmini memref.dma_start is a registered op with a fixed operand list, so it cannot carry the runtime descriptor operands the masked-DMA clamp needs (per-dim low/high are dynamic-shape values, not attributes). togsim.transfer is unregistered, so it carries them. Eliminate the dma_start intermediate: togsim.transfer flows through the whole pipeline and lowers directly to Gemmini. - New lower_transfer_to_gemmini merges decompose_transfer's <=4D handling (unit-dim collapse / >4D affine.for peel + lane-banked SRAM offset) with lower_dma_to_gemmini's CONFIG/CONFIG2/CONFIG3/[CONFIG4]/MVIN|MVOUT asm. A/B-validated (func7 + packed CONFIG rs1/rs2) against the old decompose->lower chain on add/pad/matmul/gather. - togsim.transfer gains a tag_idx operand (dram,dram_idx,sram,sram_idx,tag,tag_idx, dma_type,vst[,offset]) so the async DMA<->barrier runtime slot survives without memref.dma_start's variadic tag indices. - memref.dma_wait -> togsim.wait (togsim-level counterpart, emitted by the vcix matmul lowering, erased by the merged lowering). No memref.* DMA ops remain. - decompose_transfer dropped from PRE_OPT; loop-padding runs OPAQUELY on togsim.transfer (mlir-opt --allow-unregistered-dialect), doing only its loop-bound rounding (its DRAM-grow half is what the future masked-DMA clamp replaces). - dma_fine_grained, lower_to_vcix._DmaView, build_skeleton ported to read togsim.transfer. Functional path (Spike, timing_mode=0) validated: add, matmul (fine-grained subtile + async + togsim.wait), softmax, gather/scatter (indirect CONFIG4), constant_pad, conv (padding=0) all allclose. Known-pending: conv with padding>0 regresses -- it relied on loop-padding's DRAM-grow, which the masked-DMA clamp (next step) replaces (padding=0 passes; wrong values, not a crash).
…wait Finishes the memref.dma_start removal for the trace/timing path. build_tog's TogBuilder now decodes togsim.transfer (dram,dram_idx,sram,sram_idx,tag,tag_idx, dma_type,vst[,offset]) and togsim.wait (tag,tag_idx,num_elements) instead of memref.dma_start/dma_wait: op dispatch, _dma_start_fields (positional decode + direction from dma_kind / dma_type), the tag-user base-address scan, and the wait operand reads. TOGDMANode.tile_stride keeps reading the dram_stride attr (unchanged semantics). build_skeleton's wait marker/dispatch move to togsim.wait so barriers still become togsim.memory_barrier. dep_analysis needs no change (the decoder keeps the same field keys). Trace/timing (TOGSim, timing_mode=1) validated: add 5/5, matmul 11/11 (async DMA + tag_idx barrier pairing). No memref.* DMA ops remain in either path.
… not operand count len(operands) > 8 is ambiguous once the masked-DMA clamp adds more optional operands (low/high after offset). Gate on the `indirect` attr instead, so offset vs mask operands are disambiguated by explicit flags. Behavior-preserving: gather/add pass on both the functional (Spike) and trace (TOGSim) paths.
…gsim.wait) The trace producer now reads togsim.wait (not memref.dma_wait); update the pass docstring and comments to match. No code change.
…o CONFIG/2/3/4) Instead of packing the transfer params into four CONFIG asm instructions, build a 144-byte DMA descriptor as an 18xi64 memref.global (.data, non-constant so runtime fields can be written) and emit a single CONFIG_DESC that hands its address to the MVIN/MVOUT (which read the struct, matching the Spike-side torchsim_config_desc). This is the vehicle for the masked-DMA low/high clamp: dim_low/dim_high/fill are descriptor fields (currently defaulted, so behavior is unchanged), not new special registers. - _pack_desc encodes the fields into the struct's i64 slots (dim_size/low/high/mm_stride/ spad_stride/element_size/vlane/config_type/flags/indirect stride+esize). - descriptor globals are deduped by content; the indirect offset spad address is a runtime value, stored into slot 15 before the transfer (flags.bit0 = indirect). Requires the Spike torchsim_config_desc instruction (riscv-isa-sim). Validated functional (Spike) + trace (TOGSim): add, matmul (subtile+async), softmax, gather/scatter (indirect runtime store), constant_pad all pass.
The togsim.transfer descriptor lowering requires the spike torchsim_config_desc instruction (riscv-isa-sim PR #5, tag v1.0.4). Track the new tag so build-from-source picks up the descriptor-capable spike.
The masked-DMA clamp needs spike to read the descriptor dim_low/dim_high + desc_fill, and index_add needs the MVOUT accumulate flag. Pin v1.0.5.
…-dividing tails Drop the index_expr / indirect tile-divisibility recompile so pad_vlane_tile's lane-aligned tiles are kept instead of shrunk to a divisor of the dim. A non-dividing tile's per-tile remainder is handled by a masked-DMA clamp: - _emit_clamp emits per tile axis a dynamic [low, high) as affine.min/max of that axis' loop iv (the tile base), so the last partial tile and the pad borders fall out per iteration; low = max(0, glo - base), high = min(tile, ghi - base). The pointwise path (_masked_bounds) derives [glo, ghi) per dim (MVOUT [0, output); MVIN of a padded input [pad, pad+input_extent)); the MLIR templates (def_dma_op, e.g. gemm) derive it from the DRAM extent so a matmul tile padding past M/N/K is filled with 0 instead of MAC-ing garbage. - emit_transfer carries the (low, high) index operands + masked_axes; lower_transfer_to_gemmini writes the runtime low/high into the descriptor dim_low/dim_high through an i32-view (byte layout unchanged) + masked flag + desc_fill. One desc_ptr helper builds every descriptor: a fully static one is a deduped i64 global, while any runtime-written field (masked dim_low/dim_high and/or the indirect address, stored as two i32 words) uses a unique i32-view global, so a single descriptor carries both a masked clamp and a gather/scatter, and indirect DMAs drop their divisibility recompile too. - The box-excluded fill is the consuming reduction identity (reduction_init: sum 0, prod 1, max -inf, min +inf; 0 otherwise -- matmul MAC and stores use 0). Log-sum- exp exception (softmax / log_softmax): a sum whose node came from exp reduces exp(input), so the PRIMARY input (indexed by a reduction itervar) is filled with -inf; broadcast operands (the subtracted max) keep their finite identity, else input - (-inf) = +inf -> exp = inf -> NaN. - index_add: instead of the compute gather-add-overwrite (which loses duplicate indices landing in the same non-dividing tile), the MVOUT accumulates out[idx] += val via an accumulate flag; the DMA processes positions sequentially so duplicates accumulate correctly regardless of tile size. acc_float selects float vs integer add (element_size alone is ambiguous). With no caller left, is_dim_dividable / adjust_tile_to_divisible and the TileConstraint.must_divide_dim machinery (trim_large_tail bias, adjust loop) are removed. Fixes elementwise, constant_pad (mobilenet ragged vlane holes), reductions (sum/max/min/prod/softmax/log_softmax/var/std/layernorm/mean), gather/scatter (index_add with duplicates) and matmul/addmm on non-dividing shapes. Requires spike with the descriptor dim_low/dim_high clamp, desc_fill and the MVOUT accumulate flag.
…ne-grained)
Conv on a non-dividing reduction (K = in_ch * kh * kw, e.g. 3*7*7=147 > systolic
size) MAC-ed garbage from the tail of the K tile. Two fixes:
- def_dma_op now takes the per-dim clamp extent from the kernel's loop_extents
{loop_iv_name: loop_bound}, which each conv template declares once (the template's
ranges/itervars equivalent). The stride-match extent inference is unreliable for
conv because the gemm reads a REPACKED weight (O_C contiguous, stride 1) whose
strides do not match the declared node_layout, so the inferred extent is wrong
(e.g. O_C=64 read as 8); gemm (no loop_extents) keeps the inference. Only a dim
whose index is a single iv is clamped -- an im2col spatial axis (o_h + k_h) has no
single base and is left unclamped, which is correct because that input is
pre-padded. All four conv templates (base, mt, sb, sbs) declare loop_extents; note
the mt template fuses channel and kernel-width so its tile_k loops 0..I_C*K_W.
- dma_fine_grained now PROPAGATES the masked (and indirect / accumulate) attrs +
the low/high operands when it splits a tile into subtiles, and REMAPS each clamp
per subtile: a subtile at offset iv*sub sees high = min(sub, high - iv*sub),
low = max(0, low - iv*sub) (out-of-window subtiles get a negative high / low > sub
and Spike skips them). Previously it dropped the clamp, so the conv clamp never
reached Spike; gemm was unaffected because its DMAs do not go through this pass.
Fixes conv2d on non-dividing in_channels / kernels (test_conv2d 13/13), including
fused pointwise epilogues (conv + relu/scale/bias/sigmoid). No regression:
matmul/addmm, bmm, elementwise, reduce, softmax, index_add all pass.
A fused template epilogue (matmul/conv + relu/add/gelu/...) stores its result via store_epilogue, which -- unlike the plain store() -- emitted no masked-DMA clamp. On a non-dividing output the store then wrote the systolic pad region past the real extent, corrupting the result (fused matmul+relu at 10x10x10 was off by ~7). store_epilogue now clamps each output tile dim to [0, extent): the extent of every loop iv is ranges[k], and the tile dims are indexed by dim_aliasing (tile-dim order), so a dim whose iv extent does not divide the tile is clamped. Keyed by iv name, it is naming-agnostic -- it covers both the index-N matmul epilogue and the c0/tile_n/... conv epilogue -- and _emit_clamp skips dividing dims, so a dividing output is a no-op. Fixes fused matmul/conv epilogues on non-dividing shapes (matmul+relu/add/gelu, conv+relu/scale/bias/sigmoid, non-dividing K and non-dividing output). No regression: test_matmul, matmul_activation, addmm_residual, matmul_scalar all pass.
Hand-write a masked MVOUT togsim.transfer, run lower_transfer_to_gemmini in-process (no torch, no spike; skipped without the MLIR bindings), and compare to an embedded golden -- the exact lowered IR (MLIR SSA numbering is deterministic for a fixed input): the i32-view descriptor global (dim_size / config_type / masked flag packed), the runtime dim_low/dim_high stores at i32 idx 6/10 for masked_axes=[2], and the CONFIG_DESC + MVOUT asm. A second case adds subtile_size=[1,1,4,8] and shows the descriptor dim_size become the subtile (config) shape. Timing mode erases the transfer. Under tests/lowering/ so further per-pass goldens have a home. Regenerate the golden and review the diff on an intended lowering change.
…dims, cleanup Addresses the PR review of the masked-DMA non-dividing work: - _masked_fill_bits detected a log-sum-exp reduction with `"exp" in origin.name`, which also matches expand / expm1 / experimental -- an ordinary non-dividing sum whose origin name contains "exp" would then get a -inf tail fill instead of the sum identity 0 and reduce to -inf. Match the op target instead (_origin_is_exp: the target's overload packet name == "exp"); same fix applied to the get_padding_type twin. Unit test in tests/lowering/test_masked_fill.py pins that expand/expm1 do not match. - Remove the self._last_local_dims temporal coupling: get_dma_info now returns local_dims and load()/store() pass it to _masked_bounds explicitly, so a reordered or skipped get_dma_info can no longer feed stale clamp axes. - Drop the underscore-prefixed locals in def_dma_op / store_epilogue (non-idiomatic), trim the oversized _masked_fill_bits / _masked_bounds docstrings to the invariant plus a link to docs/design/masked-dma-descriptor.md, and clarify the bf16 dtype-table comment (present only for table totality; bf16 is not actually supported).
Dropping the tile-divisibility recompile means non-dividing shapes rely entirely on the masked-DMA clamp/fill; a gap is now a silent wrong answer, not a slow-but-safe recompile. Add tests/ops/misc/test_masked_nondividing.py pinning that path: 2-D constant pad (per-dim greedy pad recovery), sum over a broadcast operand (0 fill, guards the exp-substring bug), softmax (genuine -inf fill), non-dividing elementwise, and non-dividing gather. Register it in the CI allowlist (pytorchsim_test.yml). Also add the missing trailing newline to test_indirect_access.py.
1e43e9a to
cfcb7a9
Compare
9dbc257 to
908ea30
Compare
…asked-DMA clamp Make the MLIR backend emit valid IR under torch.compile(dynamic=True), where a loop range is a sympy symbol (e.g. s52) instead of a concrete int. - is_symbolic_dim: single predicate for "runtime dim"; the tile-fit heuristics (trim_large_tail, get_padding_ratio, make_choices) skip a symbolic dim rather than doing concrete-int arithmetic on it (kept fixed init tile; tail becomes a runtime remainder). get_mlir_shape emits a dynamic memref dim (memref<?>). - LoopLevel emits a symbolic upper bound as an index SSA (%<name>_bound); codegen_loops materializes it (memref.load + index_cast) in the prologue. - set_ranges registers each symbolic dim's runtime-extent SSA once (dyn_bound_cses), the single source of %<name>_bound (dyn_bound_name); the loop bound and the masked-DMA clamp both read that var instead of rebuilding the name at codegen time. - Masked-DMA clamp: a symbolic tile axis emits the same per-dim [low, high) clamp as a static non-dividing axis, with the runtime extent as the affine.min symbol operand -- so a non-dividing dynamic shape is handled by one path and the last partial tile falls out at runtime. - axis_split: aligned axis-split detection is symbolic-aware (+ unit test). All guards gate on is_symbolic_dim, so the static path is unchanged.
…gnostic validation, producer shape_args Drive a dynamic kernel through compile + timing + functional validation without a concrete extent. - cycle_table.pin_loops_to_one_tile: per-tile cost is shape-invariant, so sample cycles on a one-tile copy of the (symbolic) IR; extension_codecache sizes the sampling host buffers to one tile (_concretize_attrs_for_sampling). - mlir_caller_codegen: the Spike validation binary is shape-agnostic -- it loads each size arg into N_<sym> and mallocs / builds the memref descriptors from N at runtime, so one binary serves any size. extension_codecache builds and runs it for a dynamic kernel too; the per-tile spad measurement stays valid, so the spad-overflow gate runs unconditionally (static path unchanged). - lower_to_emitc: the trace producer re-sources its loop bound from shape_args, so one trace.so serves any size. build_tog/build_skeleton skip the bfs serialization that assumes a concrete loop_end. - is_dynamic_shape is detected from the size-symbol arg (MLIR_ARGS_VAR), not an IR-text sniff.
…via the attribute file The dynamic trace producer reads its loop bounds from shape_args; feed them at simulation time through the existing per-kernel attribute YAML (which already carries address_info), not a bespoke channel. - simulator: a scalar input (a dynamic size arg, e.g. s52) is not a tensor address -- collect such scalars into a shape_args sequence in the YAML, in arg order (== the producer's shape_args[k]); pass --attribute alongside --trace_so. dump_args/write_arg write the size arg's runtime value (int64) so the shape-agnostic Spike binary can load its loop bound. - main.cc: add --attribute; in the trace branch load the YAML and fill shape_args, passed to run_producer (was nullptr, 0).
A single torch.compile(dynamic=True) add compiles to one trace producer .so and is simulated at several input sizes -- the producer reads its loop bound from shape_args at runtime, so the same .so serves any size. Exercises the whole dynamic-shape pipeline end to end, including non-dividing sizes (the masked-DMA symbolic clamp handles the remainder tile).
908ea30 to
0797c45
Compare
d9131ba to
1ff2ebb
Compare
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Summary
Adds
torch.compile(dynamic=True)support on the C++ trace path. One compiledtrace producer
.soserves any input size: the producer reads its loop boundfrom
shape_argsat runtime, while the per-tile cost table (sampled once, since atile's cost is shape-invariant) keys the timing.
Verified end to end: a single dynamic
a + bruns at N=1024 (2 tiles, 183cycles) and N=2048 (4 tiles, 261 cycles) from the same
.so, the tracescaling with
shape_args.What each commit does
heuristics for a sympy dim (they only minimize a known dim's tail); all gated
on
isinstance(.., sympy.Expr), so static is unchanged.scalar-int kernel arg;
memref<?xf32>;affine.for ... to %<name>_boundwith atop-of-function
memref.load+index_castprologue.arg_attributesso thegenerated wrapper imports (the extent already arrives as a runtime arg).
binary can't be instantiated for a runtime extent.
pin_loops_to_one_tile(generalfor static + dynamic) runs gem5 sampling on a one-tile copy; the symbolic IR is
kept for the producer / cost table.
build_skeletonskips the loop_endserialization it never needed;
lower_to_emitcre-sources each still-used sizearg from
shape_args[k]viaemitc.subscript, so the producer loop readsfor (iv=0; iv<shape_args[k]; iv+=step).existing per-kernel attribute YAML (
shape_args),run_standalonepasses--attribute, andmain.ccfillsshape_argsforrun_producer.tests/ops/elementwise/test_dynamic_add.py.Known limitations / follow-ups
skipped for dynamic, so output values are zeros today; the test checks output
shape. Functional output for dynamic is a follow-up.
pin_loops_to_one_tileonce itis decoupled from
run_tog; op coverage beyond contiguous 1D add (matmul /LLVM fork and is not part of this PR.
Base: stacked on
feature/togsim-cpp-trace(#267).🤖 Generated with Claude Code